# _*_ coding:utf-8 _*_
# @Time  : 2023.07.24
# @Author: zizlee
import json
import pathlib
import time
from typing import List

import pandas as pd
from fastapi import APIRouter, Depends, Body
from pydantic import BaseModel
from response import APIResponse
from utils import datetime_utils
from utils.jwt_token import TokenPayload
from db_utils import IndexDataBaseConnection, BaseDatabaseConnection
from settings import FILE_STORAGE, STATIC_HOST


# 用户提取数据时的参数
class ExtraGroupDataModel(BaseModel):
    variety: List[str]
    groups: List[int]
    block_id: int


class ExtraGroupData:
    def __init__(self, user_id: int):
        self.user_id = user_id

    # 获取用户定制的数据指标情况
    def _query_user_flex_index(self, extra_item: ExtraGroupDataModel):
        # print(extra_item)
        # sql = """
        #     SELECT x.block_id,x.index_name As group_name,x.frequency,x.unit,
        #            x.lasted_years,x.start_date,x.end_date,x.active_expire,
        #            y.index_id,y.index_name,y.variety_en,y.block_group_id,y.plat,y.plat_id,y.db_name,y.tb_name,
        #            y.i_start,y.i_end
        #     FROM
        #         (SELECT b.id As bgid, b.block_id,b.index_name,b.frequency,b.unit,b.start_date As sd,b.end_date As ed,
        #         a.id,a.block_group_id,a.lasted_years,a.start_date,a.end_date,a.active_expire
        #         FROM datlib_flex_block_group_user As a
        #         INNER JOIN datlib_flex_block_group As b ON a.block_group_id=b.id
        #         WHERE a.user_id=%s AND a.active_expire>=CURDATE() AND b.block_id=%s AND b.is_active=1
        #               AND b.id IN %s
        #         ) As x
        #     RIGHT JOIN
        #         (SELECT m.index_id,m.index_name,m.variety_en,m.block_group_id,
        #         n.plat,n.plat_id,n.db_name,n.tb_name,n.start_date As i_start,n.end_date As i_end
        #         FROM datlib_flex_block_group_index As m
        #         INNER JOIN datlib_index As n ON m.index_id=n.id
        #         WHERE m.is_active=1 AND m.variety_en IN %s
        #         ) As y
        #     ON x.bgid=y.block_group_id
        #     WHERE y.block_group_id IN %s AND x.active_expire>=CURDATE();
        # """
        # prm = [self.user_id, extra_item.block_id, extra_item.groups, extra_item.variety, extra_item.groups]

        sql = """
            SELECT x.id,x.plat,x.plat_id,x.db_name,x.tb_name,x.start_date AS i_start,x.end_date AS i_end,x.frequency,x.unit,
                y.index_id,y.index_name,y.variety_en,y.block_group_id,y.sort_num
            FROM datlib_index AS x
            INNER JOIN
                (SELECT a.index_id,a.index_name,a.variety_en,a.block_group_id,b.is_active,b.sort_num
                FROM datlib_flex_block_group_index AS a
                INNER JOIN datlib_flex_block_group AS b ON a.block_group_id=b.id
                WHERE b.block_id=%s AND b.id IN %s AND a.variety_en IN %s)
            AS y ON x.id=y.index_id WHERE y.is_active=1 
            ORDER BY y.sort_num;
        """
        prm = [extra_item.block_id, extra_item.groups, extra_item.variety]

        with BaseDatabaseConnection() as conn:
            conn.cursor.execute(sql, prm)
            user_flex_index = list(conn.cursor.fetchall())
        return user_flex_index

    # 获取指标的数据
    def _query_index_data(self, index_list):
        # 整理共同的库的指标
        db_indexes = {}
        for iid in index_list:
            fs, fe = iid['i_start'], iid['i_end']
            if iid.get('start_date'):
                fs = iid['start_date']
            if iid.get('end_date'):
                fe = iid['end_date']
            if iid.get('lasted_years'):
                fs = datetime_utils.date_interval(iid['i_end'], year=-iid['lasted_years'], fmt='%Y-%m-%d')
                fe = iid['i_end']
            iid['fs'] = fs
            iid['fe'] = fe
            db_indexes.setdefault(iid['db_name'], []).append(iid)
        index_data_list = []
        for db_name, index_list in db_indexes.items():
            with IndexDataBaseConnection(db_name=db_name) as conn:
                for o_idx in index_list:
                    o_data = conn.query_with_interval(tb_name=o_idx['tb_name'],
                                                      start_date=o_idx['fs'], end_date=o_idx['fe'],
                                                      ascending=False)
                    index_data_list.append({**o_idx, 'data': o_data})
        return index_data_list

    def _format_index_data(self, data_list):
        ret_df = None
        table_columns = []
        for i, dt in enumerate(data_list):
            tmp = pd.DataFrame(dt['data']['rows'])
            if tmp.empty:
                continue
            value_col = f'col_{i}'
            tmp.rename(columns={'datavalue': value_col}, inplace=True)

            # for k in dt:
            #     if k != 'data':
            #         print(k, dt[k])

            table_columns.append({'prop': value_col, 'label': dt['index_name'], 'frequency': dt['frequency'], 'unit': dt['unit']})
            del tmp['id']
            if ret_df is None:
                ret_df = tmp
                continue
            ret_df = pd.merge(ret_df, tmp, how='outer', on='datadate')
        if ret_df is None:
            return ret_df, table_columns
        ret_df.sort_values(by='datadate', inplace=True, ascending=False)
        ret_df.fillna('', inplace=True)
        return ret_df, table_columns

    def extra_group_data(self, extra_item: ExtraGroupDataModel):
        # 获取用户定制的数据指标情况
        # import time
        # t1 = time.time()
        user_flex_index = self._query_user_flex_index(extra_item=extra_item)
        # print(user_flex_index)
        if len(user_flex_index) < 1:
            return APIResponse.success(data={})
        # print('_query_user_flex_index', time.time() - t1)
        # 获取指标数据
        # t2 = time.time()
        index_list_data = self._query_index_data(user_flex_index)
        # print('_query_index_data', time.time() - t2)
        # 格式化为表格形式给前端
        # t3 = time.time()
        df, columns = self._format_index_data(index_list_data)
        # print('_format_index_data', time.time() - t3)
        # t4 = time.time()
        # 缓存数据到磁盘，等待用户下载
        tmp_folder = pathlib.Path(FILE_STORAGE).joinpath('sysTemp/flexDataLib')
        if not tmp_folder.exists():
            tmp_folder.mkdir(parents=True)
        if df is None or df.empty:
            xlsx_file = tmp_folder.joinpath('00-no-data.xlsx')
            file_url = STATIC_HOST + str(xlsx_file).replace('\\', '/').replace(FILE_STORAGE, '')
            return APIResponse.success(data={'table_data': [], 'table_columns': columns, 'tmp_file': file_url})
        xlsx_file = tmp_folder.joinpath(f'{time.time_ns()}_{self.user_id}.xlsx')
        xlsx_header = {c['prop']: c['label'] for c in columns}
        xlsx_header['datadate'] = '日期'
        tmp_df = df.rename(columns=xlsx_header)
        tmp_df.to_excel(xlsx_file, index=False, engine='xlsxwriter')
        # print('cache_file', time.time() - t4)
        file_url = STATIC_HOST + str(xlsx_file).replace('\\', '/').replace(FILE_STORAGE, '')
        # t5 = time.time()
        table_data = df.to_dict(orient='records')
        # print('to_dict:', time.time() - t5)
        # import zlib
        # zlib_json = zlib.compress(str(table_data).encode('utf8'), level=5)
        # import sys
        # print(sys.getsizeof(table_data) / 1024)
        # print(sys.getsizeof(zlib_json) / 1024)
        return APIResponse.success(data={'table_data': table_data, 'table_columns': columns, 'tmp_file': file_url})


db_extra_api = APIRouter()


@db_extra_api.post('/')  # 提取数据
def extra_group_data(payload: TokenPayload = Depends(TokenPayload), extra_item: ExtraGroupDataModel = Body(...)):
    return ExtraGroupData(user_id=payload.user_id()).extra_group_data(extra_item=extra_item)
